#!/usr/bin/env python
# coding: utf-8
import pandas
import matplotlib.pyplot as plt
from PIL import Image
import os
import cv2
from tqdm import tqdm
import shutil
import pickle
classes = {
            0 : 'smoke',
            1 : 'naked',
            2 : 'mouse',
            3 : 'cat',
            4 : 'dog',
            5 : 'normal'
        }
def convert_bin_to_label(x:int):
    bin_value = bin(x)[2:]
    label = []
    for i, c in enumerate(bin_value[::-1]):
        if c == '1':
            label.append(classes[i])
    if not label:
        return 'normal'
    return ','.join(label)

if False:
    pred = pandas.read_csv('test.csv')
    file_dir = '/root/code/tianchi/data/test/picture'
    pkl_filename = 'test_imgs.pkl'
    out_dir = 'out/out_test'
    
else:
    pred = pandas.read_csv('train.csv')
    file_dir = '/root/code/tianchi/data/train'
    pkl_filename = 'train_imgs.pkl'
    out_dir = 'out/out_train'
if os.path.isfile(pkl_filename):
    print('loading pickle')
    with open(pkl_filename, 'rb') as f:
        images = pickle.load(f)
        
else:
    images =[]
    for dirpath, dirs, files in os.walk(file_dir):
        images.extend([os.path.join(dirpath, file) for file in files if file.endswith('jpg')])
    print(f'image_names\n {images}')
    images = [cv2.imread(img) for img in images]
    with open(pkl_filename, 'wb') as f:
        pickle.dump(images, f)
# print(f'images in read\n {images}')
print(f'images in pred\n {pred["filename"]}')
shutil.rmtree(out_dir, ignore_errors=True)
os.makedirs(out_dir)

    

for img, row in zip(tqdm(images), pred.itertuples(index=False)):
    filename = row.filename
    result = row.result

    # Set the title using the platform-specific method
    title = f'result: {convert_bin_to_label(int(result))}'

    # Add the title as text on top of the image
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.8
    font_thickness = 1
    text_size = cv2.getTextSize(title, font, font_scale, font_thickness)[0]
    text_position = (10, 30)  # Adjust the position as needed
    # text_position2 = (30, 30)
    border_size = 1
    img = cv2.putText(
        img,
        title,
        (10, 30),  # Adjust the position as needed
        font,
        font_scale,
        (0, 0, 0),  # Black color for the text
        font_thickness,
        cv2.LINE_AA,
    )
    img = cv2.putText(
        img,
        title,
        (10, 30),  # Adjust the position as needed
        font,
        font_scale,
        (255, 255, 255),  # White color for the border
        border_size,
        cv2.LINE_AA,
    )


    # Save the image
    output_path = os.path.join(out_dir, filename)
    cv2.imwrite(output_path, img)

print("Images saved successfully.")
